import ml_collections

import reflow.configs.default as default_config

# ! 必须定义 get_config 函数


def get_config():
    config = default_config.get_config()

    # diffusers
    config.diffusers = diffusers = ml_collections.ConfigDict()
    diffusers.ckpt_path = 'checkpoints/AltDiffusion'
    # NOTE whether load score model ckpt from pipeline_ckpt ; will be covered by score_model_ckpt
    diffusers.load_score_model = False
    diffusers.score_model_ckpt = ''

    # xlm_roberta_tokenizer, clip_tokenizer
    diffusers.tokenizer = 'xlm_roberta_tokenizer'
    # xlm_roberta_text_model, clip_text_model
    diffusers.text_encoder = 'xlm_roberta_text_model'
    diffusers.vae = 'autoencoder_kl'
    diffusers.score_model = 'unet_2d_condition_model'
    diffusers.gradient_checkpointing = False
    diffusers.use_xformers = False

    # sampling
    sampling = config.sampling
    sampling.batch_size = 5
    sampling.decode_noise = False
    sampling.decode_latent = True
    sampling.return_traj = False
    sampling.compute_loss = True
    sampling.randz0 = 'fix'
    sampling.use_ode_sampler = 'euler'
    sampling.sample_N = 10
    sampling.guidance_scale = 1.0

    sampling.method = 'rectified_flow'
    sampling.init_type = 'gaussian'
    sampling.init_noise_scale = 1.0

    # reflow
    config.reflow = reflow = ml_collections.ConfigDict()
    # NOTE: t0, t1, uniform, or an integer k > 1
    reflow.reflow_t_schedule = 'uniform'
    reflow.reflow_loss = 'l1'  # NOTE: l2, lpips, lpips+l2, msssim+l1, l1

    # data
    data = config.data
    data.eval_root = ''
    data.dl_workers = 1
    data.shuffle = False

    config.device = 'cuda:0'
    config.seed = 23

    return config
